# encoding=utf-8

import argparse
import datetime
import os
import time
import traceback

import cv2
import numpy as np

from CeilNerveformerWrapper import CeilNerveformerWrapper

ceil_nerveformer_infer = CeilNerveformerWrapper()


def run_one(file_path: str):
    try:
        t1 = time.perf_counter()
        out = ceil_nerveformer_infer.infer(file_path)
        t2 = time.perf_counter()
        print("Infer file: {}, cost: {}s".format(os.path.basename(file_path), (t2 - t1)))
        return out
    except Exception as e:
        print(traceback.format_exc())
        return None


def generate_raw_data(from_path: str, ceil_raw: str, nerve_raw: str):
    os.makedirs(ceil_raw, exist_ok=True)
    os.makedirs(nerve_raw, exist_ok=True)

    file_names = os.listdir(from_path)
    for each_name in file_names:
        abs_path = os.path.join(from_path, each_name)
        ceil_path = os.path.join(ceil_raw, each_name)
        nerve_path = os.path.join(nerve_raw, each_name)

        out = run_one(abs_path)
        if out is not None:
            ceil_mask, nerve_mask = out
            cv2.imwrite(ceil_path, ceil_mask)
            cv2.imwrite(nerve_path, nerve_mask)


def generate_color_data(from_path: str, ceil_raw: str, nerve_raw: str, color_path: str):
    os.makedirs(color_path, exist_ok=True)
    file_names = os.listdir(from_path)
    for each_name in file_names:
        t1 = time.perf_counter()
        abs_path = os.path.join(from_path, each_name)
        ceil_path = os.path.join(ceil_raw, each_name)
        nerve_path = os.path.join(nerve_raw, each_name)
        out_color_path = os.path.join(color_path, each_name)

        ori_file = cv2.imread(abs_path)
        ceil_file = cv2.imread(ceil_path)
        nerve_file = cv2.imread(nerve_path)

        no_zero_index = np.nonzero(ceil_file)
        ori_file[no_zero_index[0], no_zero_index[1], :] = \
            ori_file[no_zero_index[0], no_zero_index[1], :] * 0.5 + 0.5 * np.array([0, 212, 255])

        no_zero_index = np.nonzero(nerve_file)
        ori_file[no_zero_index[0], no_zero_index[1], :] = \
            ori_file[no_zero_index[0], no_zero_index[1], :] * 0.5 + 0.5 * np.array([255, 191, 0])

        cv2.imwrite(out_color_path, ori_file)
        t2 = time.perf_counter()
        print("Draw file: {}, cost: {}s".format(each_name, (t2 - t1)))


def run(work_dir: str):
    from_path = os.path.join(work_dir, "user_input")
    ceil_raw_path = os.path.join(work_dir, "out_raw_ceil")
    nerve_raw_path = os.path.join(work_dir, "out_raw_nerve")
    merge_path = os.path.join(work_dir, "out_merge_data")

    if not os.path.exists(from_path):
        print("user_input dir is not found in {}, program will exit".format(work_dir))
        return

    print("Start to generate model raw data in: {}, and {}".format(ceil_raw_path, nerve_raw_path))
    generate_raw_data(from_path, ceil_raw_path, nerve_raw_path)

    print("Start to generate color data in: {}".format(merge_path))
    generate_color_data(from_path, ceil_raw_path, nerve_raw_path, merge_path)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--workdir", type=str)
    args = parser.parse_args()
    if args.workdir == "" or args.workdir is None:
        print("work dir is empty")
        exit(-1)

    run(args.workdir)
